import time

import matplotlib
import numpy as np
from matplotlib import pyplot as plt

from utils import save_variable, analyze

matplotlib.use("TkAgg")
fig = plt.figure()
ax = fig.subplots()


def F(x):
    """
    目标函数

    :param x: 输入
    :return: 函数输出
    """
    return x * np.sin(10 * np.pi * x) + 1.0


class GA:
    def __init__(self, config):
        self.f = config['function']
        self.f_range = config['f_range']
        self.length = config['length']
        self.num = config['population_size']
        self.retain_rate = config['retain_rate']
        self.random_select_rate = config['random_select_rate']
        self.mutation_rate = config['mutation_rate']
        self.show_log = config['show_log']
        self.population = None
        self.init_population()

    def step(self):
        parents = self.get_selection()
        children = self.get_crossover(parents, n=self.num - parents.shape[0])
        self.population = np.concatenate([parents, children])
        self.population = self.get_mutation(self.mutation_rate)

    def search(self, num_iteration, _ax=None):
        for i in range(num_iteration):
            if self.show_log:
                print(f'epoch {i+1}/{num_iteration}:')
            self.step()
            if self.show_log:
                print(f'minimum y = {self.calc_y(self.get_result())}')
                print('=' * 50)
            if _ax is not None:
                self.plot(_ax)

        return self.get_result()

    def get_selection(self):
        sorted_population = self.population[np.argsort(self.calc_y_by_encoded_value(self.population))]
        # sorted_population = np.lexsort(self.population, keys=(self.calc_y_by_encoded_value(self.population),))
        retain_num = round(self.num * self.retain_rate)
        retained = sorted_population[:retain_num]
        other = sorted_population[retain_num:]
        random_selected = other[np.random.rand(*other.shape) < self.random_select_rate]
        selection = np.concatenate((retained, random_selected))
        if self.show_log:
            print(f"选择：{selection.shape[0]}/{self.num}存活"
                  f"（高适应性{retained.shape[0]} + 随机选择{random_selected.shape[0]}）")
        return selection

    def get_crossover(self, parents, n):
        full_mask = np.left_shift(1, self.length) - 1  # 1111111

        parents_ids = np.random.randint(0, parents.shape[0], size=(n, 2))
        couple = parents[parents_ids.flat].reshape(n, 2)

        right_mask = np.left_shift(1, np.random.randint(1, self.length, size=n)) - 1  # 0000001 - 0111111
        left_mask = full_mask - right_mask
        children = np.bitwise_or(
            np.bitwise_and(right_mask, couple[:, 0]),
            np.bitwise_and(left_mask, couple[:, 1])
        )
        if self.show_log:
            print(f"交叉：重组出{children.shape[0]}个后代")

        return children

    def get_mutation(self, mutation_rate, fix_best=True):
        whether_mutate = np.random.rand(*self.population.shape) < mutation_rate
        mutation_mask = np.left_shift(
            whether_mutate,
            np.random.randint(0, self.length, size=self.population.shape)
        )
        if fix_best:
            best_idx = self.get_best_idx()
            mutation_mask[best_idx] = 0
        # print(b)
        if self.show_log:
            print(f"突变：{whether_mutate.sum()}个个体发生突变")

        return np.bitwise_xor(self.population, mutation_mask)

    def plot(self, _ax):
        # 绘制动画
        _ax.cla()
        X = np.linspace(-1, 2, num=300)
        Y = self.calc_y(X)
        _ax.plot(X, Y)
        for each in self.population:
            x = self.decode(each)
            y = self.calc_y(x)
            _ax.plot(x, y, marker="o", markersize=5, markeredgecolor="green")
        plt.pause(0.15)

    def decode(self, encoded_value):
        decoded = encoded_value / ((1 << self.length) - 1)
        decoded = decoded * (self.f_range[1] - self.f_range[0]) + self.f_range[0]
        return decoded

    def get_result(self):
        optimal = self.population[self.get_best_idx()]
        optimal_decoded = self.decode(optimal)
        return optimal_decoded

    def get_best_idx(self):
        return np.argmin(self.calc_y_by_encoded_value(self.population))

    def calc_y(self, x):
        return self.f(x)

    def calc_y_by_encoded_value(self, encoded_value):
        return self.f(self.decode(encoded_value))

    def init_population(self):
        self.population = self.rand_population(self.length, self.num)

    @staticmethod
    def rand_population(length, n):
        population = np.zeros(n, dtype=int)
        for i in range(length):
            bit = np.random.rand(n) > 0.5
            population = np.bitwise_or(np.left_shift(population, 1), bit)
        return population


if __name__ == "__main__":
    config = {
        'function': lambda x: -F(x),
        'f_range': (-1.0, 2.0),
        'length': 22,
        'population_size': 50,
        'retain_rate': 0.1,
        'random_select_rate': 0.9,
        'mutation_rate': 0.4,
        'show_log': True
    }
    num_iteration = 100
    ga = GA(config)

    # 模拟一次并显示过程
    ga.search(num_iteration, _ax=ax)
    x_max = ga.get_result()
    print("Finish. x_max:", x_max, ", y_max:", F(x_max))
    plt.show()

    # 模拟多次测试性能
    # ga.show_log = False
    # x_list = []
    # y_list = []
    # count = 5000
    # start = time.time()
    # for i in range(count):
    #     ga.init_population()
    #     ga.search(num_iteration=num_iteration)
    #     x_max = ga.get_result()
    #     x_list.append(x_max)
    #     y_list.append(F(x_max))
    # end = time.time()
    # save_variable(x_list, "x_list_GA.pkl")
    # save_variable(y_list, "y_list_GA.pkl")
    # info = {
    #     "consume": end - start,
    #     "best_x": max(x_list, key=F),
    #     "best_y": max(y_list),
    #     "worst_y": min(y_list)
    # }
    # save_variable(info, "info_GA.pkl")
